#include "jit_compiler.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

// Forward declarations for static functions
static jit_value_t* fold_binary_constants(jit_op_t op, jit_value_t* left, jit_value_t* right);
static jit_value_t* fold_unary_constant(jit_op_t op, jit_value_t* operand);
static void mark_value_used(jit_value_t* value, jit_function_t* function);
static bool has_no_side_effects(jit_op_t op);
static void remove_instruction(jit_basic_block_t* block, jit_instruction_t* inst);
static bool is_pure_expression(jit_instruction_t* inst);
static bool instructions_equivalent(jit_instruction_t* inst1, jit_instruction_t* inst2);
static void replace_value_uses(jit_function_t* function, jit_value_t* old_value, jit_value_t* new_value);

// =============================================================================
// Optimization Pass Management
// =============================================================================

int jit_optimize_function(jit_function_t* function, uint32_t optimization_level) {
    if (!function) return -1;
    
    int result = 0;
    
    switch (optimization_level) {
        case 0:
            // No optimization
            break;
            
        case 1:
            // Basic optimizations
            result |= jit_optimize_constant_folding(function);
            result |= jit_optimize_dead_code_elimination(function);
            break;
            
        case 2:
            // Standard optimizations
            result |= jit_optimize_constant_folding(function);
            result |= jit_optimize_dead_code_elimination(function);
            result |= jit_optimize_common_subexpression_elimination(function);
            break;
            
        case 3:
            // Aggressive optimizations
            result |= jit_optimize_constant_folding(function);
            result |= jit_optimize_dead_code_elimination(function);
            result |= jit_optimize_common_subexpression_elimination(function);
            result |= jit_optimize_loop_invariant_code_motion(function);
            result |= jit_optimize_register_allocation(function);
            break;
            
        default:
            return -1;
    }
    
    return result;
}

// =============================================================================
// Constant Folding Optimization
// =============================================================================

int jit_optimize_constant_folding(jit_function_t* function) {
    if (!function) return -1;
    
    int optimizations_made = 0;
    
    jit_basic_block_t* block = function->blocks;
    while (block) {
        jit_instruction_t* inst = block->first_inst;
        
        while (inst) {
            jit_instruction_t* next_inst = inst->next;
            
            // Try to fold constants for binary operations
            if (inst->src1 && inst->src2 && 
                inst->src1->is_constant && inst->src2->is_constant) {
                
                jit_value_t* folded = fold_binary_constants(inst->opcode, 
                                                          inst->src1, inst->src2);
                if (folded) {
                    // Replace instruction with constant load
                    inst->opcode = JIT_OP_CONST;
                    inst->dest->is_constant = true;
                    inst->dest->data = folded->data;
                    inst->src1 = NULL;
                    inst->src2 = NULL;
                    
                    free(folded);
                    optimizations_made++;
                }
            }
            
            // Try to fold constants for unary operations
            else if (inst->src1 && inst->src1->is_constant) {
                jit_value_t* folded = fold_unary_constant(inst->opcode, inst->src1);
                if (folded) {
                    // Replace instruction with constant load
                    inst->opcode = JIT_OP_CONST;
                    inst->dest->is_constant = true;
                    inst->dest->data = folded->data;
                    inst->src1 = NULL;
                    
                    free(folded);
                    optimizations_made++;
                }
            }
            
            inst = next_inst;
        }
        
        block = block->next;
    }
    
    return optimizations_made;
}

// Helper function to fold binary constant operations
static jit_value_t* fold_binary_constants(jit_op_t op, jit_value_t* left, jit_value_t* right) {
    if (!left || !right || !left->is_constant || !right->is_constant) return NULL;
    
    jit_value_t* result = malloc(sizeof(jit_value_t));
    if (!result) return NULL;
    
    memset(result, 0, sizeof(jit_value_t));
    result->is_constant = true;
    
    // Handle integer arithmetic
    if (jit_type_is_integer(left->type) && jit_type_is_integer(right->type)) {
        int64_t left_val = left->data.int_val;
        int64_t right_val = right->data.int_val;
        
        switch (op) {
            case JIT_OP_ADD:
                result->type = jit_type_promote(left->type, right->type);
                result->data.int_val = left_val + right_val;
                return result;
                
            case JIT_OP_SUB:
                result->type = jit_type_promote(left->type, right->type);
                result->data.int_val = left_val - right_val;
                return result;
                
            case JIT_OP_MUL:
                result->type = jit_type_promote(left->type, right->type);
                result->data.int_val = left_val * right_val;
                return result;
                
            case JIT_OP_DIV:
                if (right_val != 0) {
                    result->type = jit_type_promote(left->type, right->type);
                    result->data.int_val = left_val / right_val;
                    return result;
                }
                break;
                
            case JIT_OP_MOD:
                if (right_val != 0) {
                    result->type = jit_type_promote(left->type, right->type);
                    result->data.int_val = left_val % right_val;
                    return result;
                }
                break;
                
            case JIT_OP_EQ:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val == right_val);
                return result;
                
            case JIT_OP_NE:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val != right_val);
                return result;
                
            case JIT_OP_LT:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val < right_val);
                return result;
                
            case JIT_OP_LE:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val <= right_val);
                return result;
                
            case JIT_OP_GT:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val > right_val);
                return result;
                
            case JIT_OP_GE:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val >= right_val);
                return result;
                
            case JIT_OP_BIT_AND:
                result->type = jit_type_promote(left->type, right->type);
                result->data.int_val = left_val & right_val;
                return result;
                
            case JIT_OP_BIT_OR:
                result->type = jit_type_promote(left->type, right->type);
                result->data.int_val = left_val | right_val;
                return result;
                
            case JIT_OP_BIT_XOR:
                result->type = jit_type_promote(left->type, right->type);
                result->data.int_val = left_val ^ right_val;
                return result;
                
            default:
                break;
        }
    }
    
    // Handle floating-point arithmetic
    if ((left->type == JIT_TYPE_FLOAT || left->type == JIT_TYPE_DOUBLE) &&
        (right->type == JIT_TYPE_FLOAT || right->type == JIT_TYPE_DOUBLE)) {
        double left_val = left->data.float_val;
        double right_val = right->data.float_val;
        
        switch (op) {
            case JIT_OP_ADD:
                result->type = JIT_TYPE_DOUBLE;
                result->data.float_val = left_val + right_val;
                return result;
                
            case JIT_OP_SUB:
                result->type = JIT_TYPE_DOUBLE;
                result->data.float_val = left_val - right_val;
                return result;
                
            case JIT_OP_MUL:
                result->type = JIT_TYPE_DOUBLE;
                result->data.float_val = left_val * right_val;
                return result;
                
            case JIT_OP_DIV:
                if (right_val != 0.0) {
                    result->type = JIT_TYPE_DOUBLE;
                    result->data.float_val = left_val / right_val;
                    return result;
                }
                break;
                
            case JIT_OP_EQ:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val == right_val);
                return result;
                
            case JIT_OP_NE:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val != right_val);
                return result;
                
            case JIT_OP_LT:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val < right_val);
                return result;
                
            case JIT_OP_LE:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val <= right_val);
                return result;
                
            case JIT_OP_GT:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val > right_val);
                return result;
                
            case JIT_OP_GE:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val >= right_val);
                return result;
                
            default:
                break;
        }
    }
    
    // Handle boolean operations
    if (left->type == JIT_TYPE_BOOL && right->type == JIT_TYPE_BOOL) {
        bool left_val = left->data.bool_val;
        bool right_val = right->data.bool_val;
        
        switch (op) {
            case JIT_OP_AND:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = left_val && right_val;
                return result;
                
            case JIT_OP_OR:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = left_val || right_val;
                return result;
                
            case JIT_OP_EQ:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val == right_val);
                return result;
                
            case JIT_OP_NE:
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = (left_val != right_val);
                return result;
                
            default:
                break;
        }
    }
    
    free(result);
    return NULL;
}

// Helper function to fold unary constant operations
static jit_value_t* fold_unary_constant(jit_op_t op, jit_value_t* operand) {
    if (!operand || !operand->is_constant) return NULL;
    
    jit_value_t* result = malloc(sizeof(jit_value_t));
    if (!result) return NULL;
    
    memset(result, 0, sizeof(jit_value_t));
    result->is_constant = true;
    
    switch (op) {
        case JIT_OP_NEG:
            if (jit_type_is_integer(operand->type)) {
                result->type = operand->type;
                result->data.int_val = -operand->data.int_val;
                return result;
            } else if (operand->type == JIT_TYPE_FLOAT || operand->type == JIT_TYPE_DOUBLE) {
                result->type = operand->type;
                result->data.float_val = -operand->data.float_val;
                return result;
            }
            break;
            
        case JIT_OP_NOT:
            if (operand->type == JIT_TYPE_BOOL) {
                result->type = JIT_TYPE_BOOL;
                result->data.bool_val = !operand->data.bool_val;
                return result;
            }
            break;
            
        case JIT_OP_BIT_NOT:
            if (jit_type_is_integer(operand->type)) {
                result->type = operand->type;
                result->data.int_val = ~operand->data.int_val;
                return result;
            }
            break;
            
        case JIT_OP_ABS:
            if (jit_type_is_integer(operand->type)) {
                result->type = operand->type;
                int64_t val = operand->data.int_val;
                result->data.int_val = val < 0 ? -val : val;
                return result;
            } else if (operand->type == JIT_TYPE_FLOAT || operand->type == JIT_TYPE_DOUBLE) {
                result->type = operand->type;
                double val = operand->data.float_val;
                result->data.float_val = val < 0.0 ? -val : val;
                return result;
            }
            break;
            
        default:
            break;
    }
    
    free(result);
    return NULL;
}

// =============================================================================
// Dead Code Elimination
// =============================================================================

int jit_optimize_dead_code_elimination(jit_function_t* function) {
    if (!function) return -1;
    
    int eliminated = 0;
    
    // First pass: mark all values as potentially dead
    for (uint32_t i = 0; i < function->value_count; i++) {
        if (function->values[i]) {
            function->values[i]->spilled = false; // Reuse spilled flag as "used" marker
        }
    }
    
    // Second pass: mark used values by walking backwards from return instructions
    jit_basic_block_t* block = function->blocks;
    while (block) {
        jit_instruction_t* inst = block->last_inst;
        
        while (inst) {
            // Mark instruction's operands as used
            if (inst->opcode == JIT_OP_RETURN) {
                if (inst->src1) mark_value_used(inst->src1, function);
            } else if (inst->dest && inst->dest->spilled) {
                // If destination is used, mark operands as used
                if (inst->src1) mark_value_used(inst->src1, function);
                if (inst->src2) mark_value_used(inst->src2, function);
                if (inst->src3) mark_value_used(inst->src3, function);
            }
            
            inst = inst->prev;
        }
        
        block = block->next;
    }
    
    // Third pass: remove instructions that produce unused values
    block = function->blocks;
    while (block) {
        jit_instruction_t* inst = block->first_inst;
        
        while (inst) {
            jit_instruction_t* next_inst = inst->next;
            
            // Remove instruction if its destination is not used
            if (inst->dest && !inst->dest->spilled && has_no_side_effects(inst->opcode)) {
                remove_instruction(block, inst);
                eliminated++;
            }
            
            inst = next_inst;
        }
        
        block = block->next;
    }
    
    return eliminated;
}

// Helper function to mark a value and its dependencies as used
static void mark_value_used(jit_value_t* value, jit_function_t* function) {
    if (!value || value->spilled) return; // Already marked
    
    value->spilled = true; // Mark as used
    
    // Find the instruction that produces this value and mark its operands
    jit_basic_block_t* block = function->blocks;
    while (block) {
        jit_instruction_t* inst = block->first_inst;
        
        while (inst) {
            if (inst->dest == value) {
                if (inst->src1) mark_value_used(inst->src1, function);
                if (inst->src2) mark_value_used(inst->src2, function);
                if (inst->src3) mark_value_used(inst->src3, function);
                return;
            }
            inst = inst->next;
        }
        
        block = block->next;
    }
}

// Helper function to check if an operation has side effects
static bool has_no_side_effects(jit_op_t op) {
    switch (op) {
        case JIT_OP_ADD:
        case JIT_OP_SUB:
        case JIT_OP_MUL:
        case JIT_OP_DIV:
        case JIT_OP_MOD:
        case JIT_OP_NEG:
        case JIT_OP_EQ:
        case JIT_OP_NE:
        case JIT_OP_LT:
        case JIT_OP_LE:
        case JIT_OP_GT:
        case JIT_OP_GE:
        case JIT_OP_AND:
        case JIT_OP_OR:
        case JIT_OP_NOT:
        case JIT_OP_BIT_AND:
        case JIT_OP_BIT_OR:
        case JIT_OP_BIT_XOR:
        case JIT_OP_BIT_NOT:
        case JIT_OP_LSHIFT:
        case JIT_OP_RSHIFT:
        case JIT_OP_ABS:
        case JIT_OP_SQRT:
        case JIT_OP_SIN:
        case JIT_OP_COS:
        case JIT_OP_CAST:
        case JIT_OP_CONST:
        case JIT_OP_PARAM:
            return true;
            
        case JIT_OP_STORE:
        case JIT_OP_CALL:
        case JIT_OP_RETURN:
            return false;
            
        default:
            return false; // Conservative assumption
    }
}

// Helper function to remove instruction from block
static void remove_instruction(jit_basic_block_t* block, jit_instruction_t* inst) {
    if (!block || !inst) return;
    
    // Update linked list pointers
    if (inst->prev) {
        inst->prev->next = inst->next;
    } else {
        block->first_inst = inst->next;
    }
    
    if (inst->next) {
        inst->next->prev = inst->prev;
    } else {
        block->last_inst = inst->prev;
    }
    
    block->inst_count--;
    
    // Free instruction
    if (inst->machine_code) {
        free(inst->machine_code);
    }
    free(inst);
}

// =============================================================================
// Common Subexpression Elimination
// =============================================================================

int jit_optimize_common_subexpression_elimination(jit_function_t* function) {
    if (!function) return -1;
    
    int eliminated = 0;
    
    // For each block, maintain a hash table of available expressions
    jit_basic_block_t* block = function->blocks;
    while (block) {
        // Simple O(n²) implementation for demonstration
        // Production implementation would use hash tables
        
        jit_instruction_t* inst1 = block->first_inst;
        while (inst1) {
            if (is_pure_expression(inst1)) {
                jit_instruction_t* inst2 = inst1->next;
                
                while (inst2) {
                    if (instructions_equivalent(inst1, inst2)) {
                        // Replace all uses of inst2->dest with inst1->dest
                        replace_value_uses(function, inst2->dest, inst1->dest);
                        
                        // Remove inst2
                        jit_instruction_t* next = inst2->next;
                        remove_instruction(block, inst2);
                        inst2 = next;
                        eliminated++;
                        continue;
                    }
                    inst2 = inst2->next;
                }
            }
            inst1 = inst1->next;
        }
        
        block = block->next;
    }
    
    return eliminated;
}

// Helper function to check if instruction is a pure expression
static bool is_pure_expression(jit_instruction_t* inst) {
    return inst && has_no_side_effects(inst->opcode) && inst->dest;
}

// Helper function to check if two instructions are equivalent
static bool instructions_equivalent(jit_instruction_t* inst1, jit_instruction_t* inst2) {
    if (!inst1 || !inst2) return false;
    if (inst1->opcode != inst2->opcode) return false;
    
    // Check operands
    if (inst1->src1 != inst2->src1) return false;
    if (inst1->src2 != inst2->src2) return false;
    if (inst1->src3 != inst2->src3) return false;
    
    return true;
}

// Helper function to replace all uses of old_value with new_value
static void replace_value_uses(jit_function_t* function, jit_value_t* old_value, jit_value_t* new_value) {
    if (!function || !old_value || !new_value) return;
    
    jit_basic_block_t* block = function->blocks;
    while (block) {
        jit_instruction_t* inst = block->first_inst;
        
        while (inst) {
            if (inst->src1 == old_value) inst->src1 = new_value;
            if (inst->src2 == old_value) inst->src2 = new_value;
            if (inst->src3 == old_value) inst->src3 = new_value;
            
            inst = inst->next;
        }
        
        block = block->next;
    }
}

// =============================================================================
// Loop Invariant Code Motion (Placeholder)
// =============================================================================

int jit_optimize_loop_invariant_code_motion(jit_function_t* function) {
    if (!function) return -1;
    
    // This is a placeholder implementation
    // Real LICM would require:
    // 1. Loop detection
    // 2. Dominance frontier analysis
    // 3. Data flow analysis
    // 4. Moving invariant computations to loop preheader
    
    return 0;
}

// =============================================================================
// Register Allocation (Placeholder)
// =============================================================================

int jit_optimize_register_allocation(jit_function_t* function) {
    if (!function) return -1;
    
    // This is a placeholder implementation
    // Real register allocation would require:
    // 1. Live range analysis
    // 2. Interference graph construction
    // 3. Graph coloring or linear scan algorithm
    // 4. Spill code insertion
    
    // For now, just assign sequential register IDs
    int reg_count = 0;
    
    for (uint32_t i = 0; i < function->value_count; i++) {
        if (function->values[i] && !function->values[i]->is_constant) {
            function->values[i]->reg_id = (reg_count % 16); // Assume 16 registers
            reg_count++;
        }
    }
    
    function->register_usage = reg_count;
    
    return 0;
}